import math
import re
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from datetime import datetime, timedelta
import json
import os
import smtplib
import threading
from docx import Document
from docx.shared import RGBColor, Pt
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT, WD_COLOR_INDEX  # 新增WD_COLOR_INDEX导入
from docx.enum.style import WD_STYLE_TYPE
from docx.oxml.ns import qn
from docx.oxml import OxmlElement
import tkinter.messagebox  # 新增导入
from ConfigManager import ConfigManager
from DocxStyleHandler import DocxStyleHandler
from JsonUtil import JsonUtil
from LLMProcess import LLMProcess
from PromptGenerator import PromptGenerator
from CustomDialog import CustomDialog
from WordProcessor import WordProcessor 
import Logger_config
from ui.ProgressBar import ProgressOverlay 
from TimePredictor import TimePredictor

# 设置日志配置
logger = Logger_config.setup_logger()

class WordOptimizer:
    def __init__(self,root): 
        self.root = root       
        self.cm = ConfigManager()
        
    def optimize(self, config: dict):
        # 入参校验
        if not self.validate_config(config):
            return
        
        file_path = config.get('file_path_optimize') # 获取输入文件路径
        output_file_path = config.get('optimize_output_file_path') # 获取输出文件路径
        
        #大模型参数获取
        model_base_url = self.cm.get_merged_config()['model']['base_url']  # 模型的基础URL
        model_api_key = self.cm.get_merged_config()['model']['api_key']  # 模型的API密钥
        model_name = self.cm.get_merged_config()['model']['model_name']  # 模型名称
        model_max_tokens = self.cm.get_merged_config()['model']['max_tokens']  # 设置模型生成的最大token数，控制生成内容的长度
        model_temperature = self.cm.get_merged_config()['model']['temperature']  # 设置模型的温度参数，控制生成内容的随机性，值越高生成的内容越随机
        model_stream = self.cm.get_merged_config()['model']['stream']  # 设置是否以流式方式获取模型生成的内容，True表示逐步返回生成结果，False表示一次性返回.

        word_style_schema = self.cm.get_merged_config()['word']['style']['schema']
        style_system_prompt = self.cm.get_merged_config()['system_prompt']['style']
        style_system_prompt = style_system_prompt.replace("{word_style_schema}",word_style_schema)

        wp = WordProcessor()
        # 一次性读取所有文本内容，用来生成样式库   
        try:         
            content = wp.read_word_document(file_path)
        except Exception as e:
            CustomDialog(self.root, "Word读取错误", "读取Word文档时出错：{str(e)}", type='info')
            return
        
        user_input_prompt = PromptGenerator.generate_prompt("optimize",config)

        userprompt = f"""
{user_input_prompt}
【文档内容】：
{content}
"""        
        prompts = [
            {"role": "system", "content": style_system_prompt},
            {"role": "user", "content": userprompt}
        ]

        logger.info(f"样式库-系统提示词长度：{len(style_system_prompt)}，内容：{style_system_prompt}")  
        logger.info(f"样式库-用户提示词长度：{len(userprompt)}，内容：{userprompt}")

        tp = TimePredictor()
        times = tp.predict_time(len(userprompt), "optimize")

         # 延迟到主线程创建进度条
        self.root.after(0, self._show_progress_bar, times)         
        
        # 在新线程中执行模型调用
        threading.Thread(
            target=self._optimize_word_file_in_thread,
            args=(model_name, model_api_key,model_base_url, model_max_tokens, model_temperature,model_stream, prompts,file_path, output_file_path,wp),
            daemon=True
        ).start()

        
    def validate_config(self, config: dict) -> bool:
        """
        入参校验
        
        参数:
            config (dict): 入参
            
        返回:
            bool: 入参是否有效            
        """
        if not config.get('file_path_optimize'):# 输入文件路径校验
            tkinter.messagebox.showwarning("警告", "请选择待优化Word文档")
            return False
        if not config.get('optimize_output_file_path'): # 输出文件路径校验
            tkinter.messagebox.showwarning("警告", "请选择输出Word文档路径")
            return False

        # 获取内容优化选项
        content_options = [
            config.get('spell_optimize', False),
            config.get('punctuation_optimize', False),
            config.get('grammar_optimize', False),
            config.get('redundancy_optimize', False),
            config.get('terminology_optimize', False),
            config.get('word_optimize', False),
            config.get('sentence_optimize', False),
            config.get('readness_optimize', False)
        ]
        
        # 获取排版优化状态和个性化提示词
        layout_enabled = config.get('enable_layout_check', False)
        custom_prompt = config.get('custom_prompt', '').strip()
        custom_prompt_placeholder = config.get('custom_prompt_placeholder', '').strip()
        
        # 判断是否为默认占位符内容        
        is_placeholder = custom_prompt == custom_prompt_placeholder
        
        # 校验三个条件是否都未满足
        if (not any(content_options) and 
            not layout_enabled and 
            (not custom_prompt or is_placeholder)):
            tkinter.messagebox.showerror(
                "错误", 
                "【内容优化】、【排版优化】、【个性化提示词】至少选择一项！"
            )
            return False

        return True
    
    def _show_progress_bar(self,times):
        steps = [("提示词生成",1), ("样式库生成",30), ("样式库写入",2),("内容及排版优化",times-1-30-2)]
        self.progress_bar = ProgressOverlay(self.root, steps)  # 保存引用
        self.progress_bar.show()
        return self.progress_bar  # 返回进度条实例
        
    def _optimize_word_file_in_thread(self, model_name, model_api_key, model_base_url, model_max_tokens, model_temperature, model_stream, prompts,file_path,output_file_path,wp:WordProcessor):
        """在子线程中执行模型调用"""
        # 调用大模型获取样式定义
        llm_result = LLMProcess.getLLMTextResponse(
            model_name, model_api_key, model_base_url,
            model_max_tokens, model_temperature, model_stream,
            prompts
        )
        logger.info(f"大模型返回样式库结果：{llm_result}")
        word_styles = JsonUtil.reshape(llm_result)

        try:
            formatted_word_styles = json.loads(word_styles)
        except json.JSONDecodeError as e:
            logger.error(f"JSON字符串不符合JSON标准。错误信息: {e}")
            logger.error(f"原始内容: {word_styles}")
            raise ValueError(f"JSON字符串不符合JSON标准。错误信息: {e}")

        dsh = DocxStyleHandler()
        doc = Document()
        dsh.create_or_update_styles(doc, formatted_word_styles) #根据大模型返回样式库，在word中创建或更新样式

        # 获取排版定义提示词
        layout_system_prompt = self.cm.get_merged_config()['system_prompt']['layout']
        
        all_results = []
        chunks = wp.read_word_doc(file_path)  # 返回 List[str]

        for i, chunk in enumerate(chunks):
            if (i + 1) % 5 == 0:
                logger.info(f"正在处理第 {i+1} 个块...")

            if len(chunks) > 1:  # 多块内容时，添加块信息
                comment = f"""
                    【文档块信息】：
                    【文档内容】共切分为{len(chunks)}个块，当前是第{i + 1}块。
                """
            else:  # 单块内容时，不添加块信息
                comment = ""

            userprompt = f"""
                【样式库】：
                {word_styles}  
                {comment}              
                【文档内容】：
                {chunk}
            """

            prompts = [
                {"role": "system", "content": layout_system_prompt},
                {"role": "user", "content": userprompt}
            ]

            llm_result = LLMProcess.getLLMTextResponse(
                model_name, model_api_key, model_base_url,
                model_max_tokens, model_temperature, model_stream,
                prompts
            )
            applied_content = JsonUtil.reshape(llm_result)

            try:
                formatted_word_content = json.loads(applied_content)
                all_results.append(formatted_word_content)
            except json.JSONDecodeError as e:
                logger.error(f"JSON字符串不符合JSON标准。错误信息: {e}")
                logger.error(f"原始内容: {applied_content}")

        # 合并所有段落为统一结构
        final_result = {
            "paragraphs": []
        }

        for result in all_results:
            final_result['paragraphs'].extend(result.get('paragraphs', []))
        logger.info(f"样式应用结果：{final_result}")

        dsh.apply_styles_and_save(output_file_path,doc,final_result)
        
        # 回到主线程隐藏进度条并弹窗
        self.root.after(0, lambda: [
            self.progress_bar.safe_hide(),
            CustomDialog(self.root, "AI优化完成", "AI优化已完成", type='info')
        ])